#pragma once

#include "EquationSolver.h"

/**
 * @file MultigridSolver.h
 * @author xingyifan
 * @date 2022-08-28 16:29
 *
 * @description: 边界条件类
 */

template <int Dim>
class BoundaryCondition
{
public:
    Real alpha_[2 * Dim];  // 狄利克雷边值系数
    Real beta_[2 * Dim];   // 诺依曼边值系数
    FuncX sigma_[2 * Dim]; // 边界条件函数

    FuncX f_;            // 离散函数
    Domain<Dim> domain_; // 离散区域

public:
    BoundaryCondition(FuncX f, Domain<Dim> &domain) : f_(f), domain_(domain) {}

    /**
     * @file EquationFiller.h
     * @author xingyifan
     * @date 2022-07-31 22:59
     *
     * @description: 填充泊松方程的边界条件函数，用于计算对应的边界条件
     */

    void fillPoisson(FuncX u, FuncX *du)
    {
        for (int d = 0; d < Dim; d++)
        {
            sigma_[2 * d] = [=](Vector x) -> Vector
            {
                return alpha_[2 * d] * Real(u(x)) - beta_[2 * d] * Real(du[d](x));
            };
            sigma_[2 * d + 1] = [=](Vector x) -> Vector
            {
                return alpha_[2 * d + 1] * Real(u(x)) + beta_[2 * d + 1] * Real(du[d](x));
            };
        }
    }
};

/**
 * @file TensorOperator.h
 * @author xingyifan
 * @date 2022-07-26 00:19
 *
 * @description: 张量/网格操作
 */

class GridOperator
{
public:
    /**
     * @file TensorOperator.h
     * @author xingyifan
     * @date 2022-07-26 00:19
     *
     * @description: 满权算子
     */

    template <int Dim>
    void FullWeighting(Grid<Dim> &tj, Grid<Dim> &t)
    {
        // 刷新指标
        tj.fresh();
        t.fresh();

        // 计算权的分布
        Tensor<1> baseWeight(3, {1, 2, 1});
        Tensor<Dim> weight = TensorOp::power<Dim>(baseWeight);

        do
        {
            // 将粗网格的指标发送到细网格上的对应位置
            tj.fine(t);

            // 不在边界上才进行加权，否则直接赋值
            if (tj.findBound() == -1)
            {
                // 获取周围一圈的值
                Tensor<Dim> value = t.cell();

                // 计算加权平均
                tj.value() = TensorOp::dot(weight, value) / pow(4, Dim);
            }
            else
            {
                tj.value() = t.value();
            }
        } while (tj++);
    }

    /**
     * @file TensorOperator.h
     * @author xingyifan
     * @date 2022-07-26 00:19
     *
     * @description: 单射算子
     */

    template <int Dim>
    void Injection(Grid<Dim> &tj, Grid<Dim> &t)
    {
        // 刷新指标
        tj.fresh();
        t.fresh();

        // 循环赋值
        do
        {
            // 将粗网格的指标发送到细网格上的对应位置
            tj.fine(t);

            // 获得细网格指标对应的值
            tj.value() = t.value();
        } while (tj++);
    }

    /**
     * @file TensorOperator.h
     * @author xingyifan
     * @date 2022-07-26 00:20
     *
     * @description: 线性插值
     */

    template <int Dim>
    void LinearProlongation(Grid<Dim> &tj, Grid<Dim> &t)
    {
        // 刷新指标
        tj.fresh();
        t.fresh();

        // 网格属性
        int n = tj.width();

        // 先插不动的点
        do
        {
            // 获得细网格上对应的指标
            t.fine(tj);
            tj.value() = t.value();
        } while (t++);

        // 刷新指标
        tj.fresh();
        t.fresh();

        // 按照维数插值
        for (int d = 0; d < Dim; d++)
        {
            // 沿此方向的步长
            int step = IntOp::power(n, d);

            do
            {
                // 当沿着此方向的坐标为偶，且不是边界点，就对左右两个点平均插值
                if (!tj.isBound(d) && tj.index(d) % 2 == 0)
                {
                    tj.value(step) = (tj.value() + tj.value(2 * step)) / 2;
                    tj.value(-step) = (tj.value() + tj.value(-2 * step)) / 2;
                }
            } while (tj++);
        }
    }

    /**
     * @file TensorOperator.h
     * @author xingyifan
     * @date 2022-07-26 00:20
     *
     * @description: 二次插值
     */

    template <int Dim>
    void QuadraticProlongation(Grid<Dim> &tj, Grid<Dim> &t)
    {
        // 刷新指标
        tj.fresh();
        t.fresh();

        // 网格属性
        int n = tj.width();

        // 先插不动的点
        do
        {
            // 获得细网格上对应的指标
            t.fine(tj);
            tj.value() = t.value();
        } while (t++);

        // 刷新指标
        tj.fresh();
        t.fresh();

        // 按照维数插值
        for (int d = 0; d < Dim; d++)
        {
            // 沿此方向的步长
            int step = IntOp::power(n, d);

            do
            {
                // 当沿着此方向的坐标为偶，且不是边界点，就对左右两个点平均插值
                if (!tj.isBound(d) && tj.index(d) % 2 == 0)
                {
                    tj.value(step) = (-tj.value(-2 * step) + tj.value() * 6 + tj.value(2 * step) * 3) / 8;
                    tj.value(-step) = (tj.value(-2 * step) * 3 + tj.value() * 6 - tj.value(2 * step)) / 8;
                }
            } while (tj++);
        }
    }
};

/**
 * @file MultigridSolver.h
 * @author xingyifan
 * @date 2022-07-31 19:19
 *
 * @description: 多重网格法
 */

template <int Dim>
class MultigridSolver
{
public:
    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-07-26 16:29
     *
     * @description: 可选不同方案和算子
     */

    MultigridSolver(BoundaryCondition<Dim> &bc, std::string method = "VC", std::string restriction = "FullWeighting", std::string prolongation = "Linear")
        : bc_(bc), omega_(2.0 / 3), method_(method), restriction_(restriction), prolongation_(prolongation) {}

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-07-28 22:58
     *
     * @description: 根据 method 选择不同方案
     */

    void solve(Tensor<Dim> &v, int nu1 = 4, int nu2 = 4, int times = 50)
    {
        Tensor<Dim> f = getf(v);
        Tensor<Dim> tv;

        // 循环求解，直到达到最大循环次数 times 或者相邻两次误差范数足够小
        do
        {
            tv = v;
            if (method_ == "VC")
            {
                VC(f, v, nu1, nu2);
            }
            if (method_ == "FMG")
            {
                FMG(f, v, nu1, nu2);
            }
            times--;
        } while (TensorOp::grid(tv - v) / TensorOp::grid(v) > 1e-6 && times > 0);
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-07-26 14:09
     *
     * @description: V 循环
     */

    void VC(Tensor<Dim> &f, Tensor<Dim> &v, int nu1, int nu2)
    {
        // 网格分割
        int n = f.width() - 1;

        // 进行 nu1 次松弛
        GS(f, v, nu1);

        // 如果不是最粗的网格
        if (n > 2)
        {
            // 网格算子
            GridOperator op;

            // 粗网格规模
            int n2 = n / 2;

            // 获得残量
            Tensor<Dim> dv = f - Laplacian(v);

            // 粗网格容器规模
            Tensor<Dim> v2(n2 + 1, 0);
            Tensor<Dim> f2(n2 + 1, 0);

            // 选择不同的约束算子
            if (restriction_ == "FullWeighting")
            {
                op.FullWeighting(f2, dv);
            }
            if (restriction_ == "Injection")
            {
                op.Injection(f2, dv);
            }

            // 在加粗网格上 V 循环
            VC(f2, v2, nu1, nu2);

            // 线性插值校正
            if (prolongation_ == "Linear")
            {
                op.LinearProlongation(dv, v2);
            }
            if (prolongation_ == "Quadratic")
            {
                op.QuadraticProlongation(dv, v2);
            }

            // 插值校正
            v = v + dv;
        }

        // 进行 nu2 次松弛
        GS(f, v, nu2);
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-07-26 16:30
     *
     * @description: FMG 循环
     * 通用的 FMG 方法只能循环一次，无论初始估计是多少，都会从零估计开始，因此精度相当低
     * 这里改成了多次循环的格式
     */

    void FMG(Tensor<Dim> &f, Tensor<Dim> &v, int nu1, int nu2)
    {
        // 网格分割
        int n = f.width() - 1;

        // 如果不是最粗的网格
        if (n > 2)
        {
            // 网格算子
            GridOperator op;

            // 粗网格规模
            int n2 = n / 2;

            // 粗网格容器规模
            Tensor<Dim> v2(n2 + 1, 0);
            Tensor<Dim> f2(n2 + 1, 0);

            // 获得残量
            Tensor<Dim> dv = f - Laplacian(v);

            // 选择不同的约束算子
            if (restriction_ == "FullWeighting")
            {
                op.FullWeighting(f2, dv);
            }
            if (restriction_ == "Injection")
            {
                op.Injection(f2, dv);
            }

            // 在粗网格上 FMG 循环
            FMG(f2, v2, nu1, nu2);

            // 线性插值校正
            if (prolongation_ == "Linear")
            {
                op.LinearProlongation(dv, v2);
            }
            if (prolongation_ == "Quadratic")
            {
                op.QuadraticProlongation(dv, v2);
            }

            // 插值校正
            v = v + dv;
        }

        // 最后做 V 循环
        VC(f, v, nu1, nu2);
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-28 12:27
     *
     * @description: Gauss-Seidel 迭代
     */

    void GS(Tensor<Dim> &f, Tensor<Dim> &v, int times)
    {
        // 网格规模
        int n = f.width();

        // 每段长度
        Real h = bc_.domain_.width() / (n - 1);

        // 加权循环
        while (times > 0)
        {
            // 遍历每一个点
            do
            {
                // 填充 ghost cell
                Vector ghost(2 * Dim, 0);

                Real center = Ghost(v, ghost);

                // 记录总和
                Real sum = 0;

                // 加权迭代
                for (int d = 0; d < Dim; d++)
                {
                    sum -= (ghost[2 * d] + ghost[2 * d + 1]) / (h * h);
                }

                v.value() = (1 - omega_) * v.value() + omega_ * (sum + f.value()) / center;
            } while (v++, f++);
            times--;
        }
    }

     /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-28 15:21
     *
     * @description: Laplacian
     */

    Tensor<Dim> Laplacian(Tensor<Dim> &v)
    {
        // 网格规模
        int n = v.width();

        // 存放结果
        Tensor<Dim> res(n, 0);

        // 每段长度
        Real h = bc_.domain_.width() / (n - 1);

        // 遍历每一个点
        do
        {
            // 填充 ghost cell
            Vector ghost(2 * Dim, 0);

            Real center = Ghost(v, ghost);

            // 记录总和
            Real sum = 0;

            // 计算乘法
            for (int d = 0; d < Dim; d++)
            {
                sum += (ghost[2 * d] + ghost[2 * d + 1]) / (h * h);
            }

            res.value() = sum + v.value() * center;
        } while (v++, res++);

        return res;
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-28 16:39
     *
     * @description: 填充 ghost cell
     */

    Real Ghost(Tensor<Dim> &v, Vector &ghost)
    {
        // 网格规模
        int n = v.width();

        // 每段长度
        Real h = bc_.domain_.width() / (n - 1);

        // 中心系数，它是 A[i,i] 元素的值
        Real center = -2 * Dim / (h * h);

        // 循环每一个方向，填充 ghost cell，同时检查是否存在狄利克雷条件
        for (int d = 0; d < Dim; d++)
        {
            // 在此方向上移动一个点的位置间隔
            int step = IntOp::power(n, d);

            // 检查是否是边界
            int sign = v.isBound(d);

            // 如果不是边界
            if (sign == 0)
            {
                // 相邻点作为 ghost cell，此时 ghost cell 真实存在
                ghost[2 * d] = v.value(-step);
                ghost[2 * d + 1] = v.value(step);
            }
            else
            {
                // 每一维上有两个边界条件 -1 表示下边界，1 表示上边界
                int ind = (sign == 1) ? 1 : 0;

                // 获得边界条件
                Real alpha = bc_.alpha_[2 * d + ind];
                Real beta = bc_.beta_[2 * d + ind];

                // 只要有一个狄利克雷条件，说明该行系数只有一个 1
                if (beta == 0)
                {
                    // 直接清空所有 ghost cell，退出循环
                    for (int i = 0; i < 2 * Dim; i++)
                    {
                        ghost[i] = 0;
                    }
                    // 中心系数为 1
                    center = 1;

                    break;
                }
                // 否则就填充 ghost cell，先获取当前点，计算边界条件
                else
                {
                    ghost[2 * d + ind] = v.value(-sign * step) - 2 * h / beta * alpha * v.value();
                    ghost[2 * d + 1 - ind] = v.value(-sign * step);
                }
            }
        }
        return center;
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-28 15:48
     *
     * @description: 填充 f
     */

    Tensor<Dim> getf(Tensor<Dim> &v)
    {
        // 系统规模
        int n = v.width() - 1;
        Tensor<Dim> f(n + 1, 0);

        // 每段长度
        Real h = bc_.domain_.width() / n;

        // 遍历每一个点填充 f
        do
        {
            // 获得当前点
            Vector np = bc_.domain_.gridPoint(f);

            // 直接获得值填充
            f.value() = Real(bc_.f_(np));

            // 在边界上，此时考虑是否存在狄利克雷条件
            if (f.findBound() >= 0)
            {
                // 循环每一个方向
                for (int d = 0; d < Dim; d++)
                {
                    // 检查是否是边界
                    int sign = f.isBound(d);

                    // 如果是边界
                    if (sign != 0)
                    {
                        // 每一维上有两个边界条件 -1 表示下边界，1 表示上边界
                        int ind = (sign == 1) ? (d * 2 + 1) : (d * 2);

                        Real beta = bc_.beta_[ind];

                        // 只要有一个狄利克雷条件，就直接取边界条件
                        if (beta == 0)
                        {
                            f.value() = Real(bc_.sigma_[ind](np));
                            break;
                        }
                        // 这里采用了一个小变化
                        /**
                         * @file MultigridSolver.h
                         * @author xingyifan
                         * @date 2022-08-28 16:24
                         *
                         * @description: 我们原本可以计算 u-1
                         * u-1 = u1 - 2h/beta (alpha u0 - sigma)
                         * 但是这种写法会使得迭代法中右端 sigma 不是任何 u 的系数，无法迭代
                         * 改写为 u-1 - 2h/beta sigma = u1 - 2h/beta alpha u0
                         * 此时右端 1, -2h/beta alpha 分别对应 u1,u0 的系数，这样就可以进行迭代
                         * 为了保持一致，原本 (u1-2u0+u-1)/h^2 = f
                         * 应当改写为 (u1-2u0+u-1-2h/beta sigma)/h^2 = f - h/(h*beta)*simga
                         * 就得到下面的式子
                         */

                        else
                        {
                            f.value() -= 2 / h / beta * Real(bc_.sigma_[ind](np));
                        }
                    }
                }
            }
        } while (f++);

        return f;
    }

public:
    Real omega_;                 // 权
    BoundaryCondition<Dim> &bc_; // 边界条件

    std::string method_;       // 求解方案
    std::string restriction_;  // 约束算子
    std::string prolongation_; // 延拓算子
};
